import os
import zipfile
import shutil


def smart_recursive_unzip(source_dir, extract_dir):
    """
    Prompt:
    用pyhon帮我实现功能：指定一个目录，遍历其中的.zip文件，如果文件名字叫"base.zip"则直接解压缩内部所有文件至当前目录（不带目录）。如果遇见其他.zip文件，则带着目录（目录名字即为.zip的文件名）解压缩至当前目录。
    注意：
    1. 指定目录中可能有多层子目录，每个子目录中都可能有.zip文件和下一层子目录，因此不能假设只有一层目录，要递归遍历。
    2. 指定一个解压目录，将所有解压出来的内容都放入这个指定解压目录，不要和原目录搞混。
    <
    智能递归解压目录下的ZIP文件

    参数:
    source_dir (str): 包含ZIP文件的源目录路径
    extract_dir (str): 解压目标目录路径
    """
    # 检查源目录是否存在
    if not os.path.exists(source_dir):
        print(f"错误：源目录 '{source_dir}' 不存在")
        return False

    # 创建解压目录（如果不存在）
    os.makedirs(extract_dir, exist_ok=True)

    # 使用os.walk递归遍历源目录及其所有子目录[1](@ref)
    for root, dirs, files in os.walk(source_dir):
        for filename in files:
            if filename.endswith('.zip'):
                zip_file_path = os.path.join(root, filename)

                # 计算相对路径，用于在解压目录中保持相同的目录结构
                rel_path = os.path.relpath(root, source_dir)
                target_base_path = os.path.join(extract_dir, rel_path)

                # 确保目标基础路径存在
                os.makedirs(target_base_path, exist_ok=True)

                print(f"处理文件: {zip_file_path}")

                try:
                    # 检查是否为有效的ZIP文件[6](@ref)
                    if not zipfile.is_zipfile(zip_file_path):
                        print(f"  警告：{zip_file_path} 不是有效的ZIP文件，跳过")
                        continue

                    with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
                        if filename == "base.zip":
                            # 对于base.zip，直接解压到当前目录（不带目录）[7](@ref)
                            zip_ref.extractall(target_base_path)
                            print(f"  已将 {filename} 内容直接解压到: {target_base_path}")
                        else:
                            # 对于其他ZIP文件，创建以文件名命名的子目录[1](@ref)
                            subdir_name = filename[:-4]  # 去掉.zip后缀
                            extract_path = os.path.join(target_base_path, subdir_name)

                            # 创建子目录
                            os.makedirs(extract_path, exist_ok=True)

                            # 解压到子目录
                            zip_ref.extractall(extract_path)
                            print(f"  已将 {filename} 解压到子目录: {extract_path}")

                except zipfile.BadZipFile:
                    print(f"  错误：{zip_file_path} 是损坏的ZIP文件")
                except Exception as e:
                    print(f"  处理 {zip_file_path} 时发生错误: {str(e)}")

    print("解压操作完成")
    return True


# 使用示例
if __name__ == "__main__":
    # 替换为您的实际路径
    source_directory = r"E:\projects_collector\appstate\pip"  # 包含ZIP文件的源目录
    extract_directory = r"E:\projects_collector\appstate\pip_unzip"  # 解压目标目录

    # 调用函数
    smart_recursive_unzip(source_directory, extract_directory)